# -*- coding: utf-8 -*-
import wx
import os
import pickle
import datetime
from configuration.ConfigData import ConfigData
from configuration.SettingWindow import SettingWindow
from configuration.ExoskeletonWindow import ExoskeletonWindow
from configuration.TrainModelWindow import TrainModelWindow
from MIController import MIController


# 初始化配置窗体
class InitWindow(wx.Frame):
    def __init__(self):
        super(InitWindow, self).__init__(None, title="主界面", size=(320, 470))
        self.SetWindowStyle(wx.DEFAULT_FRAME_STYLE & ~(
            wx.RESIZE_BORDER | wx.MAXIMIZE_BOX))
        self.config_data = ConfigData()
        self.dataServer = None
        self.init_ui()
        # Bind: 响应button事件
        self.newSubjectBtn.Bind(wx.EVT_BUTTON, self.on_new_subject)
        self.settingCueBtn.Bind(wx.EVT_BUTTON, self.on_cue_settings)
        self.settingExoBtn.Bind(wx.EVT_BUTTON, self.on_exo_settings)
        self.subjectNameCtrl.Bind(wx.EVT_CHOICE, self.on_load_param)
        self.affectedSideCtrl.Bind(wx.EVT_CHOICE, self.on_affected_side_change)
        self.saveParamBtn.Bind(wx.EVT_BUTTON, self.on_save_param)
        self.acqBtn.Bind(wx.EVT_BUTTON, self.on_graz_start)
        self.onlineBtn.Bind(wx.EVT_BUTTON, self.on_graz_start)
        self.TrainModelBtn.Bind(wx.EVT_BUTTON, self.on_train_model)

    def init_ui(self):
        self.Centre()
        self.DestroyChildren()
        # wx.Panel: 窗口的容器
        panel = wx.Panel(self)
        if not os.path.exists(self.config_data.dataSetPath):
            os.makedirs(self.config_data.dataSetPath)
        self.subjectNameList = os.listdir(self.config_data.dataSetPath)
        title_img_dir = os.path.abspath(os.getcwd()) + r'\cue_material\title_img.png'
        image = wx.Image(title_img_dir, wx.BITMAP_TYPE_PNG).ConvertToBitmap()
        title_bitmap = wx.StaticBitmap(panel, -1, image)

        # wx.FlexGridSizer: 二维网状布局(rows, cols, vgap, hgap)=>(行数, 列数, 垂直方向行间距, 水平方向列间距)
        grid_sizer1 = wx.FlexGridSizer(cols=3, vgap=5, hgap=1)
        label = wx.StaticText(panel, label="被试：")
        grid_sizer1.Add(label, 0, wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 5)
        self.subjectNameCtrl = wx.Choice(panel, name="Subject Name", choices=self.subjectNameList, size=(100, 27))
        grid_sizer1.Add(self.subjectNameCtrl, 0, wx.ALL, 5)
        self.saveParamBtn = wx.Button(panel, label="保存参数", size=(100, 27))
        grid_sizer1.Add(self.saveParamBtn, 0, wx.ALL, 5)

        label = wx.StaticText(panel, label="")
        grid_sizer1.Add(label, 0, wx.ALL, 5)
        self.acqBtn = wx.Button(panel, label="① 校准", name=self.config_data.sessionTypeSet[0], size=(100, 27))
        grid_sizer1.Add(self.acqBtn, 0, wx.ALL, 5)
        self.onlineBtn = wx.Button(panel, label="③ 在线训练", name=self.config_data.sessionTypeSet[1], size=(100, 27))
        grid_sizer1.Add(self.onlineBtn, 0, wx.ALL, 5)

        label = wx.StaticText(panel, label="训练文件：")
        grid_sizer1.Add(label, 0, wx.ALL | wx.ALIGN_CENTER_HORIZONTAL, 5)
        self.trainDataNum = wx.SpinCtrl(panel, value='1', min=1, max=6, size=(80, 27))
        grid_sizer1.Add(self.trainDataNum, 0, wx.ALIGN_CENTER_HORIZONTAL, 5)
        self.TrainModelBtn = wx.Button(panel, label="② 模型训练", size=(100, 27))
        grid_sizer1.Add(self.TrainModelBtn, 0, wx.ALIGN_CENTER_HORIZONTAL, 5)

        sbox = wx.StaticBox(panel, -1, label=u'参数配置')
        sbsizer = wx.StaticBoxSizer(sbox, wx.VERTICAL)

        grid_sizer2 = wx.FlexGridSizer(cols=3, vgap=10, hgap=1)
        self.newSubjectBtn = wx.Button(panel, label="新建被试", size=(120, 27))
        grid_sizer2.Add(self.newSubjectBtn, 0, wx.ALL, 5)
        label = wx.StaticText(panel, label="患侧：")
        grid_sizer2.Add(label, 0, wx.ALL | wx.ALIGN_CENTER_VERTICAL, 5)
        self.affectedSideCtrl = wx.Choice(panel, name="Affected Side", choices=['左', '右'], size=(70, 27))
        self.affectedSideCtrl.SetSelection(0)
        grid_sizer2.Add(self.affectedSideCtrl, 0, wx.ALL, 5)

        grid_sizer3 = wx.FlexGridSizer(cols=3, vgap=10, hgap=1)
        self.settingExoBtn = wx.Button(panel, label="外骨骼配置", size=(120, 27))
        grid_sizer3.Add(self.settingExoBtn, 0, wx.ALL, 5)
        self.settingCueBtn = wx.Button(panel, label="提示界面配置", size=(120, 27))
        # self.settingCueBtn.Disable()  # 禁用btn
        grid_sizer3.Add(self.settingCueBtn, 0, wx.ALL, 5)

        sbsizer.Add(grid_sizer2, proportion=0, flag=wx.ALL, border=5)
        sbsizer.Add(grid_sizer3, proportion=0, flag=wx.ALL, border=5)

        self.statusBar = self.CreateStatusBar()  # 状态栏
        self.statusBar.SetStatusText(u'……')

        gridSizer = wx.FlexGridSizer(cols=1, vgap=1, hgap=1)
        gridSizer.Add(title_bitmap, 0, wx.ALL, 5)
        gridSizer.Add(grid_sizer1, 0, wx.ALL, 5)
        gridSizer.Add(sbsizer, 0, wx.ALIGN_CENTER_HORIZONTAL, 5)

        panel.SetSizerAndFit(gridSizer)
        panel.Center()
        self.Fit()

    def on_new_subject(self, event):
        # 新建被试名文件夹
        new_subject_dlg = wx.TextEntryDialog(self, '输入新被试名：', '新建被试')
        if new_subject_dlg.ShowModal() == wx.ID_OK:
            new_subject_name = new_subject_dlg.GetValue()
            new_subject_path = self.config_data.dataSetPath + r'/' + new_subject_name
            if not os.path.exists(new_subject_path):
                os.makedirs(new_subject_path)
                self.subjectNameList.append(new_subject_name)
                self.subjectNameCtrl.SetItems(self.subjectNameList)
                self.subjectNameCtrl.SetSelection(self.subjectNameCtrl.FindString(new_subject_name))
        new_subject_dlg.Destroy()

    def on_cue_settings(self, event):
        setting_win = SettingWindow(self, "自定义提示界面参数")
        setting_win.ShowModal()
        self.config_data = setting_win.get_value()

    def on_exo_settings(self, event):
        exo_window = ExoskeletonWindow(self, "外骨骼配置")
        exo_window.ShowModal()
        self.config_data = exo_window.get_value()

    def on_affected_side_change(self, event):
        affected_side = self.affectedSideCtrl.GetString(self.affectedSideCtrl.GetCurrentSelection())
        self.config_data.set_affected_side(affected_side)

    def on_load_param(self, event):
        subject_name = self.subjectNameCtrl.GetString(self.subjectNameCtrl.GetCurrentSelection())
        if subject_name:
            file_path = self.config_data.dataSetPath + '/' + subject_name + '/' + subject_name + '_Param.pkl'
            if os.path.exists(file_path):
                with open(file_path, 'rb') as file:
                    self.config_data = pickle.loads(file.read())
                self.statusBar.SetStatusText(r'参数载入成功')
                if self.config_data.affectedSide == 'left':
                    self.affectedSideCtrl.SetSelection(self.affectedSideCtrl.FindString('左'))
                else:
                    self.affectedSideCtrl.SetSelection(self.affectedSideCtrl.FindString('右'))
            else:
                self.statusBar.SetStatusText(r'未保存参数文件')
        else:
            self.statusBar.SetStatusText(r'未选择被试')

    def on_save_param(self, event):
        subjectName = self.subjectNameCtrl.GetString(self.subjectNameCtrl.GetCurrentSelection())
        if subjectName:
            self.config_data.trainModelPath = self.config_data.dataSetPath + '/' + subjectName + '/' + subjectName + '_Model.pkl'
            filePath = self.config_data.dataSetPath + '/' + subjectName + '/' + subjectName + '_Param.pkl'
            output = open(filePath, 'wb')
            output.write(pickle.dumps(self.config_data))
            output.close()
            self.statusBar.SetStatusText(r'参数文件保存成功')
        else:
            self.statusBar.SetStatusText(r'未选择被试')

    def on_graz_start(self, event):
        subject_name = self.subjectNameCtrl.GetString(self.subjectNameCtrl.GetCurrentSelection())
        if subject_name:
            subject_path = self.config_data.dataSetPath + '/' + subject_name
            date = datetime.datetime.now().strftime('%Y%m%d')
            self.config_data.saveDataPath = subject_path + '/' + subject_name + '_' + date
            if not os.path.exists(self.config_data.saveDataPath):  # 创建 姓名_日期 文件夹
                os.makedirs(self.config_data.saveDataPath)
            sessionType = event.GetEventObject().GetName()
            if sessionType == self.config_data.sessionTypeSet[0]:
                tag = "校准"
                is_online = False
            else:
                self.config_data.trainModelPath = self.config_data.dataSetPath + '/' + subject_name + '/' + subject_name + '_Model.pkl'
                if os.path.exists(self.config_data.trainModelPath):
                    tag = "训练"
                    is_online = True
                else:
                    self.statusBar.SetStatusText(r'未找到训练模型')
                    return
            msg = "是否开始【" + tag + "任务】?"
            style = wx.OK | wx.CANCEL | wx.CENTRE
            msgbox = wx.MessageDialog(self, msg, tag+"任务开始", style)
            if msgbox.ShowModal() == wx.ID_OK:
                controller = MIController(self.config_data, is_online)
                controller.start()
            else:
                return
        else:
            self.statusBar.SetStatusText(r'未选择被试')

    def on_train_model(self, event):
        subject_name = self.subjectNameCtrl.GetString(self.subjectNameCtrl.GetCurrentSelection())
        if subject_name:
            subject_path = self.config_data.dataSetPath + '/' + subject_name
            self.config_data.trainModelPath = subject_path + '/' + subject_name + '_Model.pkl'
            train_model_win = TrainModelWindow(self, "模型训练")
            train_model_win.ShowModal()
        else:
            self.statusBar.SetStatusText(r'未选择被试')
